# ==============================
# _*_ coding:utf-8 _*_
# _*_ by:zizle     _*_
# _*_ at:20230421  _*_
# ==============================
from typing import Union
import datetime

HOLIDAYS = {
    2022: [],
    2023: [
        '20230102', '20230123', '20230124', '20230125', '20230126', '20230127', '20230405',
        '20230501', '20230502', '20230503', '20230622', '20230623', '20230929', '20231002',
        '20231003', '20231004', '20231005', '20231006'
    ],
    2024: [
        '20240101'
    ]

}


def now_yyyymmdd():
    return datetime.datetime.now().strftime('%Y%m%d')


def now_yyyy_mm_dd():
    return datetime.datetime.now().strftime('%Y-%m-%d')


def now_yyyy_mm_dd_hhmmss(link: str = ' ', tsep: str = ''):
    return datetime.datetime.now().strftime('%Y-%m-%d{}%H{}%M{}%S'.format(link, tsep, tsep))


def now_format(to_format: str):
    return datetime.datetime.now().strftime(to_format)


def to_yyyymmdd(obj: Union[datetime.date,datetime.datetime]):
    return obj.strftime('%Y%m%d')


def to_yyyy_mm_dd(obj: Union[datetime.date,datetime.datetime]):
    return obj.strftime('%Y-%m-%d')


def to_yyyy_mm_dd_hhmmss(obj: Union[datetime.date, datetime.datetime], link: str = ' ', tsep: str = ''):
    return obj.strftime('%Y-%m-%d{}%H{}%M{}%S'.format(link, tsep, tsep))


def to_custom(obj: Union[datetime.date, datetime.datetime], to_format: str = '%Y-%m-%d %H:%M:%S'):
    return obj.strftime(to_format)


def auth_yyyymmdd(s: str, obj: bool = False, tf='%Y%m%d'):
    try:
        od = datetime.datetime.strptime(str(s), '%Y%m%d')
    except ValueError:
        return None
    return od if obj else od.strftime(tf)


def auth_yyyy_mm_dd(s: str = None, obj: bool = False):
    if not s:
        return None
    try:
        od = datetime.datetime.strptime(s, '%Y-%m-%d')
    except ValueError:
        return None
    return od if obj else to_yyyy_mm_dd(od)


def auth_yyyy_mm_dd_hhmmss(s: str, obj: bool = False, link: str = ' ', tsep: str = ''):
    try:
        od = datetime.datetime.strptime(s, '%Y-%m-%d{}%H{}%M{}%S'.format(link, tsep, tsep))
    except ValueError:
        return None
    return od if obj else to_yyyy_mm_dd_hhmmss(od, tsep=tsep)


def auth_custom(s: str, to_format: str, obj: bool = False, link: str = ' ', tsep: str = ''):
    try:
        od = datetime.datetime.strptime(s, '%Y-%m-%d{}%H{}%M{}%S'.format(link, tsep, tsep))
    except ValueError:
        return None
    return od if obj else to_custom(od, to_format)


def translate_date_string(s: str, b_format: str, a_format: str):
    try:
        od = datetime.datetime.strptime(s, b_format)
        after = od.strftime(a_format)
    except ValueError:
        return None
    return after


def calculate_seconds_diff(t1: Union[str, datetime.datetime], t2: datetime.datetime = None, link: str = ' '):
    if isinstance(t1, str):
        t1 = auth_yyyy_mm_dd_hhmmss(t1, obj=True, link=link)
    if not t1:
        raise ValueError('can not parse t1 to datetime.')
    if not t2:
        t2 = auth_yyyy_mm_dd_hhmmss(now_yyyy_mm_dd_hhmmss(), obj=True)
    return (t2-t1).seconds


def interval_days(dt: Union[str, datetime.datetime, datetime.date], days: int, link='', obj=True):
    if isinstance(dt, str):
        new_dt = auth_yyyy_mm_dd_hhmmss(dt, obj=True, link=link)
        if not new_dt:
            new_dt = auth_yyyy_mm_dd(dt, obj=True)
    else:
        new_dt = dt
    assert new_dt
    return new_dt + datetime.timedelta(days=days) if obj else to_yyyy_mm_dd(new_dt + datetime.timedelta(days=days))


def date_interval(base_date_string: str, year: int = 0, month: int = 0, days: int = 0, fmt: str = None):
    base_date = auth_yyyy_mm_dd(base_date_string, obj=True)
    if not base_date:
        return base_date_string
    t_year = base_date.year + year
    t_month = base_date.month
    t_day = base_date.day
    t_date = datetime.datetime(year=t_year, month=t_month, day=t_day)
    month_day = month * 31
    delta_days = month_day + days
    t_date = t_date + datetime.timedelta(days=delta_days)
    return t_date.strftime(fmt) if fmt else t_date


def is_weekday_or_holiday(dt: Union[str, datetime.datetime, datetime.date] = None, link=''):
    new_dt = None
    if isinstance(dt, str):
        new_dt = auth_yyyy_mm_dd_hhmmss(dt, obj=True, link=link)
        if not new_dt:
            new_dt = auth_yyyymmdd(dt, obj=True)
    if not new_dt:
        new_dt = datetime.date.today()
    if new_dt.weekday() >= 5:
        return True
    return to_yyyymmdd(new_dt) in HOLIDAYS.get(new_dt.year, [])


def is_out_of_exchange(dt: Union[str, datetime.datetime, datetime.date] = None, link=''):
    weekday_or_holiday = is_weekday_or_holiday(dt, link)
    if weekday_or_holiday:
        return True
    now_str = now_yyyy_mm_dd_hhmmss(tsep=':').split(' ')[1]
    return '02:00:00' < now_str < '09:00:00' or '16:00:00' <= now_str < '21:00:00'


# 生成一年内的所有日期
def generate_dates_of_year(year: int = 2020):
    dates_list = []
    start = datetime.date(year, 1, 1)
    while start <= datetime.date(year, 12, 31):
        dates_list.append(start.strftime('%m-%d'))
        start += datetime.timedelta(days=1)
    return dates_list


def get_exchange_days(start: Union[datetime.date, datetime.datetime], end: Union[datetime.date, datetime.datetime]):  # 获取交易日
    days = []
    while start <= end:
        if start.weekday() <= 4 and start.strftime('%Y-%m-%d') not in HOLIDAYS.get(start.year, []):
            days.append(start.strftime('%Y%m%d'))
        start = start + datetime.timedelta(days=1)
    return days

